﻿using System;
using System.Collections.Generic;
using System.IO;

namespace ArgusLib.Drawing.Formats
{
	public class PNG
	{
		static readonly byte[] Signature = new byte[] { (byte)137, (byte)80, (byte)78, (byte)71, (byte)13, (byte)10, (byte)26, (byte)10 };
		static readonly ChunkType IEND = new ChunkType("IEND");

		Stream stream;
		CRC32 crcCalc;
		Dictionary<string, long> dictChunks;
		//ChunkType chunkType;
		//long chunkOffset;
		//int chunkLength;

		private PNG(Stream stream, CRC32 crc)
		{
			this.stream = stream;
			this.crcCalc = crc;

			this.dictChunks = new Dictionary<string, long>();
			//this.chunkOffset = -1;
			//this.chunkLength = 0;

			long iendOffset = this.SearchChunk(PNG.IEND);
			if (iendOffset < 0)
				throw new InvalidDataException("PNG Stream is corrupt.");
		}

		public static PNG Open(Stream stream)
		{
			if (stream.CanRead == false)
				return null;
			if (stream.CanSeek == false)
				return null;

			for (int i = 0; i < PNG.Signature.Length; i++)
			{
				if (stream.ReadByte() != PNG.Signature[i])
					return null;
			}

			CRC32 crc = CRC32.CreateTable();
			return new PNG(stream, crc);
		}

		//#region ChunkStreaming
		//public static PNG OpenChunk(Stream stream, ChunkType chunkType)
		//{
		//	PNG png = PNG.Open(stream);
		//	png.chunkOffset = png.SearchChunk(chunkType);
		//	png.chunkType = chunkType;
		//	if (png.chunkOffset > -1)
		//	{
		//		png.chunkOffset += 8;
		//		stream.Seek(png.chunkOffset, SeekOrigin.Begin);
		//		png.chunkLength = PNG.ReadInt32(stream) + 4;
		//	}
		//	return png;
		//}

		//private void Create

		//public override bool CanRead
		//{
		//	get
		//	{
		//		if (this.chunkOffset < 0)
		//			return false;
		//		return this.stream.CanRead;
		//	}
		//}
		//public override bool CanSeek { get { return this.stream.CanSeek; } }
		//public override bool CanTimeout { get { return this.stream.CanTimeout; } }
		//public override bool CanWrite { get { return this.stream.CanWrite; } }
		//public override long Length { get { return this.chunkLength; } }
		//public override long Position
		//{
		//	get
		//	{
		//		if (this.chunkOffset < 0)
		//			return 0;
		//		return this.stream.Position - this.chunkOffset;
		//	}
		//	set
		//	{
		//		if (value != 0 && this.chunkOffset < 0)
		//			throw new EndOfStreamException();
		//		this.stream.Position = value + this.chunkOffset;
		//	}
		//}

		//public override int Read(byte[] buffer, int offset, int count)
		//{
		//	if (this.chunkOffset < 0)
		//		return 0;

		//	int length = (int)Math.Min(this.chunkLength+this.chunkOffset-this.stream.Position, offset + count);
		//	return this.stream.Read(buffer, offset, length - offset);
		//}

		//public override int ReadByte()
		//{
		//	if (this.chunkOffset < 0)
		//		return -1;

		//	if (this.stream.Position < this.chunkOffset + this.chunkLength)
		//		return base.ReadByte();
		//	else
		//		return -1;
		//}

		//public override long Seek(long offset, SeekOrigin origin)
		//{
		//	if (this.chunkOffset < 0)
		//		return 0;

		//	if (origin == SeekOrigin.Begin)
		//	{
		//		if (offset > this.chunkLength)
		//			offset = this.chunkLength;
		//		this.stream.Seek(this.chunkOffset + offset, SeekOrigin.Begin);
		//	}
		//	else if (origin == SeekOrigin.Current)
		//	{
		//		long pos = this.stream.Position - this.chunkOffset;
		//		if (pos + offset > this.chunkLength)
		//			offset = this.chunkLength - pos;
		//		this.stream.Seek(offset, SeekOrigin.Current);
		//	}
		//	else if (origin == SeekOrigin.End)
		//	{
		//		if (offset > this.chunkLength)
		//			offset = this.chunkLength;
		//		this.stream.Seek(this.chunkOffset + this.chunkLength - offset, SeekOrigin.Begin);
		//	}
		//	return this.stream.Position - this.chunkOffset;
		//}
		//#endregion

		public byte[] ReadChunk(ChunkType chunkType, out bool isCrcValid)
		{
			isCrcValid = false;
			long offset = this.SearchChunk(chunkType);
			if (offset < 0)
				return null;

			this.stream.Seek(offset, SeekOrigin.Begin);
			int length = PNG.ReadInt32(this.stream);
			if (length < 1)
				return null;
			this.stream.Seek(4, SeekOrigin.Current);
			byte[] data = new byte[length];
			this.stream.Read(data, 0, length);
			int crc = PNG.ReadInt32(this.stream);

			this.crcCalc.Clear();
			this.crcCalc.Add(chunkType.TypeCode);
			this.crcCalc.Add(data);

			isCrcValid = this.crcCalc.CurrentCRC == crc;
			return data;
		}

		public void WriteChunk(ChunkType chunkType, byte[] data)
		{
			long iendOffset = this.SearchChunk(PNG.IEND);
			if (iendOffset < 0)
				throw new InvalidDataException("PNG Stream is corrupt.");

			long offset = this.SearchChunk(chunkType);
			if (offset < 0)
			{
				this.stream.Seek(iendOffset, SeekOrigin.Begin);

				this.crcCalc.Clear();
				this.crcCalc.Add(chunkType.TypeCode);
				this.crcCalc.Add(data);
				PNG.WriteInt32(this.stream, data.Length);
				PNG.WriteChunkType(this.stream, chunkType);
				this.stream.Write(data, 0, data.Length);
				PNG.WriteInt32(this.stream, this.crcCalc.CurrentCRC);

				// IEND
				PNG.WriteInt32(this.stream, 0);
				PNG.WriteChunkType(this.stream, PNG.IEND);
				int crc = this.crcCalc.Calculate(PNG.IEND.TypeCode);
				PNG.WriteInt32(this.stream, crc);
			}
			else
			{
				this.stream.Seek(offset, SeekOrigin.Begin);
				int length = PNG.ReadInt32(this.stream) + 12;
				byte[] buffer = new byte[length];
				while (true)
				{
					int count = (int)Math.Min(length, iendOffset - (offset + 2 * length));
					this.stream.Seek(offset + length, SeekOrigin.Begin);
					this.stream.Read(buffer, 0, count);
					this.stream.Seek(offset, SeekOrigin.Begin);
					this.stream.Write(buffer, 0, count);
					offset += count;
					if (count < length)
						break;
				}

				this.crcCalc.Clear();
				this.crcCalc.Add(chunkType.TypeCode);
				this.crcCalc.Add(data);
				PNG.WriteInt32(this.stream, data.Length);
				PNG.WriteChunkType(this.stream, chunkType);
				this.stream.Write(data, 0, data.Length);
				PNG.WriteInt32(this.stream, this.crcCalc.CurrentCRC);

				// IEND
				PNG.WriteInt32(this.stream, 0);
				PNG.WriteChunkType(this.stream, PNG.IEND);
				int crc = this.crcCalc.Calculate(PNG.IEND.TypeCode);
				PNG.WriteInt32(this.stream, crc);

				this.stream.Flush();
				this.stream.SetLength(this.stream.Position);
			}
		}

		private long SearchChunk(ChunkType chunkType)
		{
			long offset;
			if (this.dictChunks.TryGetValue(chunkType.Name, out offset) == true)
			{
				this.stream.Seek(offset, SeekOrigin.Begin);
				int length = PNG.ReadInt32(this.stream);
				if (length > -1)
				{
					ChunkType cType = PNG.ReadChunkType(this.stream);
					if (cType == chunkType)
						return offset;
				}
			}

			this.stream.Seek(PNG.Signature.Length, SeekOrigin.Begin);
			while (true)
			{
				offset = this.stream.Position;
				int length = PNG.ReadInt32(this.stream);
				if (length < 0)
				{
					if (this.dictChunks.ContainsKey(chunkType.Name) == true)
						this.dictChunks.Remove(chunkType.Name);
					return -1;
				}
				ChunkType cType = PNG.ReadChunkType(this.stream);

				if (chunkType == cType)
				{
					string chunkName = chunkType.Name;
					if (this.dictChunks.ContainsKey(chunkName) == false)
						this.dictChunks.Add(chunkType.Name, offset);
					else 
						this.dictChunks[chunkName] = offset;
					return offset;
				}

				// Skip data + 4 Bytes CRC32
				this.stream.Seek(length+4, SeekOrigin.Current);
			}
		}

		private static int ReadInt32(Stream stream)
		{
			byte[] length = new byte[4];
			int read = stream.Read(length, 0, 4);
			if (read < 4)
				return -1;
			return
				(length[0] << 24) +
				(length[1] << 16) +
				(length[2] << 8) +
				length[3];
		}

		private static void WriteInt32(Stream stream, int i)
		{
			byte[] data = new byte[4];
			data[0] = (byte)(i >> 24);
			data[1] = (byte)(i >> 16);
			data[2] = (byte)(i >> 8);
			data[3] = (byte)(i);
			stream.Write(data, 0, 4);
		}

		private static ChunkType ReadChunkType(Stream stream)
		{
			byte[] typeCode = new byte[4];
			stream.Read(typeCode, 0, 4);
			return new ChunkType(typeCode);
		}

		private static void WriteChunkType(Stream stream, ChunkType chunkType)
		{
			stream.Write(chunkType.TypeCode, 0, 4);
		}

		public class ChunkType
		{
			public byte[] TypeCode { get; private set; }
			public string Name
			{
				get
				{
					return new string(new char[] { (char)this.TypeCode[0], (char)this.TypeCode[1], (char)this.TypeCode[2], (char)this.TypeCode[3] });
				}
			}

			public ChunkType(byte[] typeCode)
			{
				if (typeCode.Length != 4)
					throw new ArgumentException("The length of the bytearray must be 4.", "typeCode");

				for (int i = 0; i < 4; i++)
				{
					int code = (int)typeCode[i];
					if (!((code > 64 && code < 91) || (code > 96 && code < 123)))
						throw new ArgumentException(string.Format("Invalid character at index {0}.", i), "typeCode");
				}
				this.TypeCode = typeCode;
			}

			public ChunkType(char a, char b, char c, char d)
				: this(new byte[] { (byte)a, (byte)b, (byte)c, (byte)d })
			{ }

			public ChunkType(string typeCode)
				: this(typeCode[0], typeCode[1], typeCode[2], typeCode[3])
			{ }

			public override bool Equals(object obj)
			{
				ChunkType chunkType = obj as ChunkType;
				if (chunkType == null)
					return false;
				
				for (int i = 0; i < this.TypeCode.Length; i++)
				{
					if (this.TypeCode[i] != chunkType.TypeCode[i])
						return false;
				}
				return true;
			}

			public static bool operator ==(ChunkType chunkTypeA, ChunkType chunkTypeB)
			{
				if (object.Equals(chunkTypeA, null) == true)
					return object.Equals(chunkTypeB, null);
				return chunkTypeA.Equals(chunkTypeB);
			}

			public static bool operator !=(ChunkType chunkTypeA, ChunkType chunkTypeB)
			{
				if (object.Equals(chunkTypeA, null) == true)
					return !object.Equals(chunkTypeB, null);
				return !chunkTypeA.Equals(chunkTypeB);
			}
		}
	}
}
